Fix Norm Layers, Again #1509

DhairyaLGandhi · 2021-02-16T09:05:23Z

This is a successor to #1495

PR Checklist

Tests are added
Entry in NEWS.md
Documentation, if applicable
Final review from @dhairyagandhi96 (for API changes).

Todo:

Go through the tests, we don't want to use the autodiff test file as it is, doing more straightforward checks is.better.
Remove hasaffine
remove testmode! - use zygote instead and have an equivalent forward pass function available a la dropout
determine forward pass at compile time

test/cuda/layers.jl

CarloLucibello · 2021-02-17T10:59:31Z

remove testmode! - use zygote instead and have an equivalent forward pass function available a la dropout

this TODO has nothing to do with this PR.

darsnack · 2021-02-17T13:30:51Z

And the last time it was dicussed there was not a complete consensus. I think it is better not to put this TODO on the v0.12 path.

DhairyaLGandhi · 2021-02-20T14:41:29Z

Still not perfect, but better than master. We would need to figure removing the unions as well.

This:

julia> bn = BatchNorm(10)
BatchNorm(10, identity, affine = )

julia> x = rand(Float32, 3,3,3,10);

julia> @code_warntype bn(x)
Variables
  BN::BatchNorm{typeof(identity),Array{Float32,1},Float32,Array{Float32,1}}
  x::Array{Float32,4}
  #334::Flux.var"#334#335"{Array{Float32,4},Int64}
  @_4::Int64
  N::Int64
  reduce_dims::Array{Int64,1}
  affine_shape::Union{Nothing, NTuple{4,Int64}}
  ts::Union{Nothing, Bool}
  μ::Array{Float32,4}
  σ²::Array{Float32,4}
  @_11::Union{Nothing, NTuple{4,Int64}}
  @_12::Union{Nothing, Bool}

Body::Array{Float32,4}
1 ──       Core.NewvarNode(:(#334))
│          Core.NewvarNode(:(@_4))
│          Core.NewvarNode(:(reduce_dims))
│          Core.NewvarNode(:(affine_shape))
│          Core.NewvarNode(:(ts))
│          Core.NewvarNode(:(μ))
│          Core.NewvarNode(:(σ²))
│          (N = Flux.ndims(x))
[...]

Master:

julia> bn = BatchNorm(10)
BatchNorm(10)

julia> x = rand(Float32, 3,3,10,10);

julia> bn(x);

julia> @code_warntype bn(x)
Variables
  BN::BatchNorm{typeof(identity),Array{Float32,1},Float32,Array{Float32,1}}
  x::Array{Float32,4}
  #239::Flux.var"#239#240"{Array{Float32,4},Int64}
  N::Int64
  reduce_dims::Array{Int64,1}
  affine_shape::NTuple{4,Int64}

Body::Any
1 ─       Core.NewvarNode(:(#239))
│         Core.NewvarNode(:(N))
│         Core.NewvarNode(:(reduce_dims))
│         Core.NewvarNode(:(affine_shape))
│   %5  = Flux.ndims(x)::Core.Compiler.Const(4, false)
│   %6  = (%5 - 1)::Core.Compiler.Const(3, false)
│   %7  = Flux.size(x, %6)::Int64
│   %8  = Base.getproperty(BN, :chs)::Int64
│   %9  = (%7 == %8)::Bool

darsnack · 2021-03-03T13:41:43Z

Based on our discussion yesterday, this PR looks ready except the plans for active. Can we agree to save those for another release, since master currently makes no changes to the active/trainmode!/testmode! API?

Other than that looks like we only need tests and docs.

CarloLucibello · 2021-03-05T20:41:00Z

src/layers/normalise.jl

+  β = initβ(chs)
+  γ = initγ(chs)


why we should have params that we don't use in practice when affine=false?

Other thought on this: I know we rejected Affine(BatchNorm(...)) because it is not ergonomically nice. But what about having a field in BatchNorm called affine that is a function. When the affine kwarg is true, then affine = Dense(...) but when it is false, affine = identity? In a Chain, the user still see just BatchNorm and bn.affine will show a Dense which the user will intuitively understand. We also don't end up storing data we don't need, and trainable will automatically be empty.

Kind of like you'd have for an activation function? That does sound pretty appealing, would it make sense to call that field transform or something then?

Yeah exactly. transform would certainly capture the field accurately, but affine might be better here cause people will understand it.

well that had been a consideration. Its a bit awkward to also do γ = reshape(l.γ, affine_shape) but I think we can have broadcasting take care of the shape as needed by l.λ.(γ .* x̂ .+ β)

CarloLucibello · 2021-03-05T20:46:11Z

src/layers/normalise.jl

            affine, track_stats, nothing)
 end

 @functor BatchNorm
-trainable(bn::BatchNorm) = hasaffine(bn) ? (bn.β, bn.γ) : ()
+# trainable(bn::BatchNorm) = hasaffine(bn) ? (bn.β, bn.γ) : ()


if you remove this μ, σ² will end up in the params

src/layers/normalise.jl

test/cuda/layers.jl

CarloLucibello · 2021-03-05T20:55:28Z

test/cuda/layers.jl


+ pixelshuffle = [PixelShuffle]
+ gpu_gradtest("PixelShuffle 2d", pixelshuffle, rand(Float32, 3, 4, 18, 3), 3)
+ gpu_gradtest("PixelShuffle 1d", pixelshuffle, rand(Float32, 3, 18, 3), 3)


there is a weird one space indent here

I think the indent is still here?

CarloLucibello · 2021-03-07T09:45:14Z

I'll recap what this PR when ready should not change with respect to master:

the constructor interface BatchNorm(3, affine=true, track_stats=false) and its default values
the behaviour in the forward pass
the layers' parameters:

julia> params(BatchNorm(2))
Params([Float32[0.0, 0.0], Float32[1.0, 1.0]])

julia> params(BatchNorm(2, affine=false))
Params([])

So, since there won't be any API changes with respect to master (unless I'm missing something), we can consider this non-blocking for v0.12 release

Co-authored-by: Carlo Lucibello <[email protected]>

DhairyaLGandhi · 2021-03-07T10:39:45Z

How does one change the behaviour of a layer at different points of training? Say I don't want to perform affine transforms initially, but want to later. This params change would hinder that without adding extra boilerplate code on the user end and potentially extra copies slowing down the entire process unnecessarily. This breaks the contract of eager execution.

How do I also restart training from a given point if I don't serialise the state of the tracking variables at such a time? They would not be trained either way, so the user choice is respected.

CarloLucibello · 2021-03-07T11:45:57Z

How does one change the behaviour of a layer at different points of training? Say I don't want to perform affine transforms initially, but want to later. This params change would hinder that without adding extra boilerplate code on the user end and potentially extra copies slowing down the entire process unnecessarily. This breaks the contract of eager execution.

It is not clear what you say in the last two sentences, but the thing is that params function should return a meaningful result, that is whatever are meant to be a layer's trainable parameters. If you declare m = BatchNorm(2, affine=false) then params(m) should be empty, anything else would be surprising and would make params meaningless.

The use case you mention can be simply achieved by creating the layer affine=true and then freezing the parameters in any
part of the training you want the same way you would do it with any set of parameters.

How do I also restart training from a given point if I don't serialise the state of the tracking variables at such a time? They would not be trained either way, so the user choice is respected.

Again, params(m) should return trainable parameters not anything in a module. Serialization has nothing to do with params.

PyTorch handles the separation between trainable parameters and the rest of the internal state very neatly:

In [15]: import torch.nn as nn

In [16]: m = nn.BatchNorm2d(2, affine=True)

In [17]: list(m.parameters())
Out[17]: 
[Parameter containing:
 tensor([1., 1.], requires_grad=True),
 Parameter containing:
 tensor([0., 0.], requires_grad=True)]

In [18]: m = nn.BatchNorm2d(2, affine=False)

In [19]: list(m.parameters())
Out[19]: []

In [20]: m = nn.BatchNorm2d(2, track_running_stats=True)

In [21]: list(m.buffers())
Out[21]: [tensor([0., 0.]), tensor([1., 1.]), tensor(0)]

In [22]: m = nn.BatchNorm2d(2, track_running_stats=False)

In [23]: list(m.buffers())
Out[23]: []

darsnack

There's still some issues related to changes to the tests which aren't clear.

src/cuda/cudnn.jl

src/layers/normalise.jl

darsnack · 2021-06-24T14:38:32Z

test/layers/normalisation.jl

  testmode!(m, true)
-  y = evalwgrad(m, x) # should override istraining
-  @test count(a->a==0, y) == 0
+  y = m(x) # should override istraining


This still holds

darsnack · 2021-06-24T14:39:00Z

test/layers/normalisation.jl

    y = m(x)
-    @test isapprox(y, sigmoid.((x .- m.μ) ./ sqrt.(m.σ² .+ m.ϵ)), atol = 1.0e-7)
+    @test_broken isapprox(y, mean(sigmoid.((x .- m.μ) ./ sqrt.(m.σ² .+ m.ϵ)), dims = 1), atol = 1.0e-7)


test/layers/normalisation.jl

darsnack · 2021-06-24T14:41:06Z

test/layers/normalisation.jl

+  # @test length(Flux.params(InstanceNorm(10))) == 0
+  # @test length(Flux.params(InstanceNorm(10, affine = true))) == 2
+  # @test length(Flux.params(InstanceNorm(10, affine =false))) == 0


Should be changed back

We aren't guaranteeing this.

A test isn't a guarantee of behavior. But this is the current behavior so it should be tested. If two weeks from now, someone borks trainable, then this test would accurately identify the issue which is why it exists.

CarloLucibello · 2021-06-24T15:21:28Z

src/layers/normalise.jl

+    μnew, σ²new = track_stats(x, (l.μ, l.σ²), (μ,σ²),
+                              l.momentum, reduce_dims = nc.dims)


these can go in the zygote block below

I think this is okay since track_stats has @nograd

in which case the @nograd on track_stats can be removed

CarloLucibello · 2021-06-24T15:32:18Z

test/layers/normalisation.jl

  end

-  let m = trainmode!(BatchNorm(2)), x = reshape(Float32.(1:6), 3, 2, 1)


these tests shouldn't be changed

CarloLucibello · 2021-06-24T15:33:28Z

test/layers/normalisation.jl

-  let m = BatchNorm(32), x = randn(Float32, 416, 416, 32, 1);
-    m(x)
-    @test (@allocated m(x)) <  100_000_000
-  end
+  # let m = BatchNorm(32), x = randn(Float32, 416, 416, 32, 1);
+  #   m(x)
+  #   @test (@allocated m(x)) <  100_000_000
+  # end


probably this was here for a reason, could try to check with git blame when it was introduced

Blame leads back to #586. Personally I think it's a little weird to just test allocations for BatchNorm, but on the other hand this is 4 lines that run quickly and don't add any real maintenance overhead. My vote is either uncomment it or get rid of it entirely. Leaving stuff around but commented without a little explainer comment raises more questions than it answers :P

test/layers/normalisation.jl

CarloLucibello · 2021-06-24T15:36:32Z

test/layers/normalisation.jl

-    y = m(x)
-    μ = mean(x, dims=1)
-    σ² = var(x, dims=1, corrected=false)
+    y, back = pullback((m,x) -> m(x), m, x)


when internal changes to a layer previous test shouldn't be modified so that one can be sure that previous behavior is preserved. Normalization layer behave differently in gradient contexts, this test was testing the behaviour in a non-gradient context

CarloLucibello · 2021-06-24T15:37:58Z

there are too many changes in tests, they should be mostly reverted, otherwise it's hard to check that the revised internals don't introduce new (or old) bugs

darsnack · 2021-06-29T12:37:51Z

src/layers/normalise.jl

+    Zygote.ignore() do
+      l.μ = reshape(μnew, :)
+      l.σ² = reshape(σ²new, :)
    end


Why not put this in track_stats?

darsnack · 2021-06-29T12:42:46Z

src/layers/normalise.jl

-  sz = sz isa Integer ? (sz,) : sz
-  diag = affine ? Diagonal(sz...) : nothing
-  return LayerNorm(λ, diag, ϵ, sz, affine)
+function LayerNorm(sz, λ = identity; affine = Diagonal(sz...), ϵ = 1f-5)


This isn't quite what I was suggesting. I didn't mean change the affine keyword to be non-boolean. LayerNorm was fine as is. I meant storing affine as a function in BatchNorm, etc. similar to LayerNorm. The API would still be a boolean.

I get that, but for the future, we would likely not want to limit kwargs to Boolean, since while it is symmetric with other libraries, it is somewhat limiting what can be expressed. Specifically a boolean allows to express two states, while an anonymous function can do more than that. Since our implementation isn't limiting what can be done with the layer, we should pass those freedoms to the end user as well. This means, we can have a compatibility layer to support booleans, but in times when it isn't worth it to write ones own layer or tweak with certain settings that aren't directly expressed with the Boolean, it's handy to pass in an anonymous function instead.

We wouldn't want the Boolean to be the only api that users can interact with. In fact, over use of it can add bloat to the code base.

I agree completely, but the default constructor provides that flexibility? I also agree on the overuse of booleans, but for the norm layers 99.99% of users want an affine transformation (not some other function) and they just want to turn it on or off. Here, a boolean makes sense for the sake of convenience. We can still allow the advanced use case with the default constructor.

The default constructor is much less ergonomic and harder to use. Not to mention requiring the user to know way more details about Flux's innards that it's untenable and a poor use of time, not to mention making the user code fragile to updates to Flux.

The fact that other libraries push boolean kwargs is why I proposed having some kind of hybrid constructor that can accept either.

Sure, a hybrid constructor is fine by me

DhairyaLGandhi · 2021-08-12T05:38:39Z

bumping

ToucheSir

Seems like there are still a handful of comments here, so resolving those should clear up any blockers.

ToucheSir · 2021-08-12T18:00:12Z

src/layers/normalise.jl

-  diag = affine ? Diagonal(sz...) : nothing
-  return LayerNorm(λ, diag, ϵ, sz, affine)
+function LayerNorm(sz, λ = identity; affine = Diagonal(sz...), ϵ = 1f-5)
+  # diag = affine ? Diagonal(sz...) : identity


Per the discussion above (hybrid constructor), I believe this line needs to be brought back.

ToucheSir · 2021-08-12T18:02:11Z

src/layers/normalise.jl

+struct NormConfig{A,T}
+  dims::Vector{Int}
+end
+
+NormConfig(affine, track_stats, reduce_dims) = NormConfig{affine, track_stats}(reduce_dims)
+
+getaffine(nc::NormConfig{true}, sz_x; dims = length(sz_x) - 1) =
+  ntuple(i -> i in dims ? sz_x[i] : 1, length(sz_x))
+
+getaffine(nc::NormConfig{false}, args...; kwargs...) = ()
+
 # For InstanceNorm, GroupNorm, and BatchNorm.
 # Compute the statistics on the slices specified by reduce_dims.
 # reduce_dims=[1,...,N-2,N] for BatchNorm
 # reduce_dims=[1,...,N-2] for InstanceNorm and GroupNorm
-function _norm_layer_forward(l, x::AbstractArray{T,N}; reduce_dims, affine_shape) where {T, N}
-  if !_isactive(l) && l.track_stats # testmode with tracked stats
+function norm_forward(l, x::AbstractArray{T,N}, nc::NormConfig{A, true}) where {A, T, N}


Now that all this functionality has been pulled out, we should consider moving it into NNlib next. Getting rid of https://github.com/FluxML/Flux.jl/tree/master/src/cuda (which only contains a batchnorm override) would be excellent and save us quite a few headaches.

CarloLucibello · 2023-02-05T15:25:30Z

closing as stale

Dhairya Gandhi added 5 commits February 16, 2021 14:34

clean up the history

2cd59b4

bn on cudnn

0ef033e

rm extra utils

1e67700

use simpler test suite first

0dbeb39

git fixes

04ca0a8

CarloLucibello reviewed Feb 17, 2021

View reviewed changes

test/cuda/layers.jl Outdated Show resolved Hide resolved

Dhairya Gandhi added 4 commits February 19, 2021 04:28

refactor norm_forward

31f7000

simplify tests

262112a

typose

c61457c

check reduce for batch

0432b82

Dhairya Gandhi added 4 commits February 21, 2021 21:44

use normconfig

77a8e87

use normconfig for other layers

d0961a7

typo

b091a80

backwards

680e64f

DhairyaLGandhi mentioned this pull request Feb 22, 2021

Deprecate active field, and auto keyword #1301

Closed

4 tasks

first pass

4a26559

CarloLucibello reviewed Mar 5, 2021

View reviewed changes

src/layers/normalise.jl Show resolved Hide resolved

CarloLucibello reviewed Mar 5, 2021

View reviewed changes

src/layers/normalise.jl Outdated Show resolved Hide resolved

CarloLucibello reviewed Mar 5, 2021

View reviewed changes

test/cuda/layers.jl Outdated Show resolved Hide resolved

CarloLucibello reviewed Mar 5, 2021

View reviewed changes

Update src/layers/normalise.jl

1c38029

Co-authored-by: Carlo Lucibello <[email protected]>

darsnack requested changes Jun 24, 2021

View reviewed changes

DhairyaLGandhi added 3 commits June 24, 2021 20:26

space cleanups + show

0d4605d

add layer norm show methods

36084e5

whitespace

3c6f1ce

CarloLucibello reviewed Jun 24, 2021

View reviewed changes

test/layers/normalisation.jl Show resolved Hide resolved

CarloLucibello reviewed Jun 24, 2021

View reviewed changes

DhairyaLGandhi added 2 commits June 25, 2021 11:49

change some tests

8f6de19

use affine as function

c525d4f

darsnack reviewed Jun 29, 2021

View reviewed changes

ericphanson mentioned this pull request Jul 6, 2021

Workaround Flux#1027 beacon-biosignals/LegolasFlux.jl#4

Merged

ToucheSir mentioned this pull request Jul 17, 2021

deprecate DepthwiseConv once we have groups in standard conv #1667

Closed

rebase

aa39039

ToucheSir reviewed Aug 12, 2021

View reviewed changes

ToucheSir mentioned this pull request Aug 22, 2021

Frest start : Next steps FluxML/ONNX.jl#49

Open

DhairyaLGandhi mentioned this pull request Sep 30, 2021

Normalization Layers not interating well with destructure/restructure #1727

Closed

pxl-th mentioned this pull request Feb 3, 2022

Fix type-stability for normalization layers #1856

Merged

4 tasks

DhairyaLGandhi mentioned this pull request Jul 27, 2022

Deprecate Flux.Optimisers and implicit parameters in favour of Optimisers.jl and explicit parameters #1986

Closed

8 tasks

CarloLucibello closed this Feb 5, 2023

chengchingwen deleted the dg/instance2 branch February 5, 2023 17:55

chengchingwen restored the dg/instance2 branch February 5, 2023 18:04

CarloLucibello deleted the dg/instance2 branch March 22, 2023 09:12

		μnew, σ²new = track_stats(x, (l.μ, l.σ²), (μ,σ²),
		l.momentum, reduce_dims = nc.dims)

		end

		let m = trainmode!(BatchNorm(2)), x = reshape(Float32.(1:6), 3, 2, 1)

Fix Norm Layers, Again #1509

Fix Norm Layers, Again #1509

Conversation

DhairyaLGandhi commented Feb 16, 2021 • edited Loading

PR Checklist

CarloLucibello commented Feb 17, 2021 • edited Loading

darsnack commented Feb 17, 2021

DhairyaLGandhi commented Feb 20, 2021

darsnack commented Mar 3, 2021 • edited Loading

CarloLucibello Mar 5, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CarloLucibello commented Mar 7, 2021

DhairyaLGandhi commented Mar 7, 2021

CarloLucibello commented Mar 7, 2021

darsnack left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CarloLucibello commented Jun 24, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DhairyaLGandhi commented Aug 12, 2021

ToucheSir left a comment

Choose a reason for hiding this comment

ToucheSir Aug 12, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CarloLucibello commented Feb 5, 2023

DhairyaLGandhi commented Feb 16, 2021 •

edited

Loading

CarloLucibello commented Feb 17, 2021 •

edited

Loading

darsnack commented Mar 3, 2021 •

edited

Loading

CarloLucibello Mar 5, 2021 •

edited

Loading

ToucheSir Aug 12, 2021 •

edited

Loading